#include <dicom/storage/encrypt.h>
#include <dicom/storage/decrypt.h>
#include <gwbase/base/util.h>
#include <gwbase/base/test_harness.h>
#include <gwbase/base/test_util.h>
#include <gwbase/encrypt/compress_helper.h>
#include <surbasic/AESHelper.h>
#include <gwbase/encrypt/digest.h>

using namespace dicom;

const std::string key =
    "defc580a5dccfd05a7f2b5433fea3283a5b962627189b17ec3a6ae5239462c17";

class EncryptorTest {};

void CheckSurbasicN(size_t n, bool random, unsigned char byte) {
  Status s;
  const std::string src_file = StringPrintf("test_%ld.dat", n);
  const std::string expect_plain = StringPrintf("test.expect_plain_%ld.dat", n);
  const std::string out_file = StringPrintf("test.out_%ld.dat", n);
  Random rand;
  CreateTestFile(src_file, n, random, byte);

  std::string md5;
  std::string sha1;

  SDBasic::code::SD_CODE c = SDBasic::aes::AESHelper::CompressAndEncFile(
      SDBasic::bytearray::ByteArray::fromHexStr(key), src_file, out_file, &md5,
      &sha1);
  EXPECT_TRUE(c == SDBasic::code::SD_SUCCESSED);

  c = SDBasic::aes::AESHelper::DecAndUncompressFile(
      SDBasic::bytearray::ByteArray::fromHexStr(key), out_file, expect_plain);
  EXPECT_TRUE(c == SDBasic::code::SD_SUCCESSED);
  EXPECT_TRUE(CheckSameFile(expect_plain, src_file));

  LOG_INFO << "SIZE: " << n << " PASSED";
}

void CheckEncryptN(size_t n, bool random, unsigned char byte) {
  Status s;
  const std::string src_file = StringPrintf("test_%ld.dat", n);
  const std::string expect = StringPrintf("test.expect_%ld.dat", n);
  const std::string expect_plain = StringPrintf("test.expect_plain_%ld.dat", n);
  const std::string out_file = StringPrintf("test.out_%ld.dat", n);
  Random rand;
  CreateTestFile(src_file, n, random, byte);

  std::string md5;
  std::string sha1;

  SDBasic::code::SD_CODE c = SDBasic::aes::AESHelper::CompressAndEncFile(
      SDBasic::bytearray::ByteArray::fromHexStr(key), src_file, expect, &md5,
      &sha1);
  EXPECT_TRUE(c == SDBasic::code::SD_SUCCESSED);

  c = SDBasic::aes::AESHelper::DecAndUncompressFile(
      SDBasic::bytearray::ByteArray::fromHexStr(key), expect, expect_plain);
  EXPECT_TRUE(c == SDBasic::code::SD_SUCCESSED);
  EXPECT_TRUE(CheckSameFile(expect_plain, src_file));

  HashEntry hash1;
  s = dicom::ZipAndEncryptFile(key, src_file, out_file, &hash1);
  EXPECT_TRUE(s.ok());

  EXPECT_TRUE(CheckSameFile(expect_plain, src_file));

  LOG_INFO << "SIZE: " << n << " PASSED";
}

void CheckEncryptDecryptN(size_t n, bool random, unsigned char byte) {
  Status s;
  const std::string src_file = StringPrintf("test_%ld.dat", n);
  const std::string expect_plain = StringPrintf("test.expect_%ld.dat", n);
  const std::string out_file = StringPrintf("test.out_%ld.dat", n);
  Random rand;
  CreateTestFile(src_file, n, random, byte);

  HashEntry hash1;
  s = dicom::ZipAndEncryptFile(key, src_file, out_file, &hash1);
  EXPECT_TRUE(s.ok());

  s = dicom::DecryptAndUnzipFile(key, out_file, expect_plain);
  EXPECT_TRUE(s.ok());

  EXPECT_TRUE(CheckSameFile(expect_plain, src_file));

  LOG_INFO << "SIZE: " << n << " PASSED";
}

void CheckEncryptDecryptHashN(size_t n, bool random, unsigned char byte) {
  Status s;
  const std::string src_file = StringPrintf("test_%ld.dat", n);
  const std::string expect_plain = StringPrintf("test.expect_%ld.dat", n);
  const std::string out_file = StringPrintf("test.out_%ld.dat", n);
  Random rand;
  CreateTestFile(src_file, n, random, byte);

  HashEntry hash1;
  s = dicom::ZipAndEncryptFile(key, src_file, out_file, &hash1);
  EXPECT_TRUE(s.ok());

  HashEntry hash2;
  s = dicom::DecryptAndUnzipFileAndHash(key, out_file, expect_plain, &hash2);
  EXPECT_TRUE(s.ok());

  EXPECT_TRUE(CheckSameFile(expect_plain, src_file));
  EXPECT_TRUE(HashEntryEqual(hash1, hash2));

  LOG_INFO << "SIZE: " << n << " PASSED";
}

TEST(EncryptorTest, SurbasicAllZeroTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckSurbasicN(i, false, '\0');
  }

  CheckSurbasicN(1000, false, '\0');
  CheckSurbasicN(1024, false, '\0');
  CheckSurbasicN(1024 * 1024, false, '\0');
}

TEST(EncryptorTest, SurbasicRandomTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckSurbasicN(i, true, '\0');
  }

  CheckSurbasicN(1000, true, '\0');
  CheckSurbasicN(1024, true, '\0');
  CheckSurbasicN(1024 * 1024, true, '\0');
}

TEST(EncryptorTest, WriterAllZeroTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptN(i, false, '\0');
  }

  CheckEncryptN(1000, false, '\0');
  CheckEncryptN(1024, false, '\0');
  CheckEncryptN(1024 * 1024, false, '\0');
}

TEST(EncryptorTest, WriterRandomTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptN(i, true, '\0');
  }

  CheckEncryptN(1000, true, '\0');
  CheckEncryptN(1024, true, '\0');
  CheckEncryptN(1024 * 1024, true, '\0');
}

TEST(EncryptorTest, ReaderAllZeroFileTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptDecryptN(i, false, '\0');
  }

  CheckEncryptDecryptN(1000, false, '\0');
  CheckEncryptDecryptN(1024, false, '\0');
  CheckEncryptDecryptN(1024 * 1024, false, '\0');
}

TEST(EncryptorTest, ReaderRandomFileTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptDecryptN(i, true, '\0');
  }

  CheckEncryptDecryptN(1000, true, '\0');
  CheckEncryptDecryptN(1024, true, '\0');
  CheckEncryptDecryptN(1024 * 1024, true, '\0');
}

TEST(EncryptorTest, ReaderAllZeroFileHashTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptDecryptHashN(i, false, '\0');
  }

  CheckEncryptDecryptHashN(1000, false, '\0');
  CheckEncryptDecryptHashN(1024, false, '\0');
  CheckEncryptDecryptHashN(1024 * 1024, false, '\0');
}

TEST(EncryptorTest, ReaderRandomFileHashTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptDecryptHashN(i, true, '\0');
  }

  CheckEncryptDecryptHashN(1000, true, '\0');
  CheckEncryptDecryptHashN(1024, true, '\0');
  CheckEncryptDecryptHashN(1024 * 1024, true, '\0');
}

int main() { RunAllTests(); }
